import traceback

from pyasn1.codec.der import encoder, decoder
from pyasn1_modules import rfc2315, rfc5280, rfc3279, rfc5208, rfc4055, rfc2459  # 引入 rfc4055
from pyasn1.type import univ, namedtype, namedval, tag, constraint, useful
import base64
import hashlib
import os
import datetime
import random

from issuer_and_serial_number import build_issuer_name


def generate_rsa_key_pair(key_size=2048):
    """生成 RSA 密钥对（未修改）"""
    modulus = os.urandom(key_size // 8)
    public_exponent = b'\x01\x00\x01'  # 65537
    private_exponent = os.urandom(key_size // 8)
    p = os.urandom(key_size // 16)
    q = os.urandom(key_size // 16)
    dp = os.urandom(key_size // 16)
    dq = os.urandom(key_size // 16)
    inverse_q = os.urandom(key_size // 16)

    public_key_info = rfc5280.SubjectPublicKeyInfo()
    algorithm = rfc5280.AlgorithmIdentifier()
    algorithm['algorithm'] = rfc3279.rsaEncryption

    rsa_public_key = univ.Sequence()
    rsa_public_key.setComponentByPosition(0, univ.Integer(int.from_bytes(modulus, byteorder='big')))
    rsa_public_key.setComponentByPosition(1, univ.Integer(int.from_bytes(public_exponent, byteorder='big')))

    bit_string = univ.BitString.fromOctetString(encoder.encode(rsa_public_key))
    public_key_info['algorithm'] = algorithm
    public_key_info['subjectPublicKey'] = bit_string

    private_key_info = rfc5208.PrivateKeyInfo()
    private_key_info['version'] = univ.Integer(0)
    private_key_info['privateKeyAlgorithm'] = algorithm

    rsa_private_key = univ.Sequence()
    rsa_private_key.setComponentByPosition(0, univ.Integer(0))
    rsa_private_key.setComponentByPosition(1, univ.Integer(int.from_bytes(modulus, byteorder='big')))
    rsa_private_key.setComponentByPosition(2, univ.Integer(int.from_bytes(public_exponent, byteorder='big')))
    rsa_private_key.setComponentByPosition(3, univ.Integer(int.from_bytes(private_exponent, byteorder='big')))
    rsa_private_key.setComponentByPosition(4, univ.Integer(int.from_bytes(p, byteorder='big')))
    rsa_private_key.setComponentByPosition(5, univ.Integer(int.from_bytes(q, byteorder='big')))
    rsa_private_key.setComponentByPosition(6, univ.Integer(int.from_bytes(dp, byteorder='big')))
    rsa_private_key.setComponentByPosition(7, univ.Integer(int.from_bytes(dq, byteorder='big')))
    rsa_private_key.setComponentByPosition(8, univ.Integer(int.from_bytes(inverse_q, byteorder='big')))

    octet_string = univ.OctetString(encoder.encode(rsa_private_key))
    private_key_info['privateKey'] = octet_string

    public_key_der = encoder.encode(public_key_info)
    private_key_der = encoder.encode(private_key_info)

    return public_key_der, private_key_der


def generate_certificate(public_key_der, issuer=None, validity_days=365):
    """生成 X.509 证书（修复版本字段设置问题）"""
    public_key_info, _ = decoder.decode(public_key_der, asn1Spec=rfc5280.SubjectPublicKeyInfo())
    serial_number = random.randint(1, 2 ** 64 - 1)

    # 创建 TBSCertificate 结构
    tbs_certificate = rfc5280.TBSCertificate()

    # 设置版本 (v3) - 使用正确的方式设置版本字段
    # version = rfc5280.Version(2)
    print(type(tbs_certificate['version']))
    tbs_certificate['version'] = 'v1'
    # tbs_certificate.setComponentByName('version', 'v1')
    print(tbs_certificate.prettyPrint())
    tbs_certificate['version'] = univ.Integer(1).subtype(
        explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatSimple, 0)
    )
    print(tbs_certificate.prettyPrint())

    # 设置序列号
    tbs_certificate.setComponentByName('serialNumber', univ.Integer(serial_number))
    # tbs_certificate['serialNumber'] = univ.Integer(serial_number)

    # 设置签名算法
    signature_algorithm = rfc5280.AlgorithmIdentifier()
    signature_algorithm['algorithm'] = univ.ObjectIdentifier('1.2.840.113549.1.1.11')
    tbs_certificate.setComponentByName('signature', signature_algorithm)

    issuer = build_issuer_name(
        country='CN',
        organization='Example Company',
        organizational_unit='IT Department',
        common_name='Example CA'
    )

    tbs_certificate.setComponentByName('issuer', issuer)

    # 设置有效期
    validity = rfc5280.Validity()

    # 生效时间
    now = datetime.datetime.utcnow()

    not_before_time = rfc5280.Time()
    not_before_time['utcTime'] = useful.UTCTime(now.strftime('%y%m%d%H%M%SZ'))
    # 失效时间
    expiry = now + datetime.timedelta(days=validity_days)
    not_after_time = rfc5280.Time()
    not_after_time['utcTime'] = useful.UTCTime(expiry.strftime('%y%m%d%H%M%SZ'))

    # 设置 Time 类型
    validity.setComponentByName('notBefore', not_before_time)
    validity.setComponentByName('notAfter', not_after_time)

    tbs_certificate.setComponentByName('validity', validity)

    # 设置主题 (与颁发者相同，因为是自签名)
    tbs_certificate.setComponentByName('subject', issuer)

    # 设置公钥
    tbs_certificate.setComponentByName('subjectPublicKeyInfo', public_key_info)

    # 创建证书
    certificate = rfc5280.Certificate()
    certificate.setComponentByName('tbsCertificate', tbs_certificate)

    # 设置签名算法
    certificate.setComponentByName('signatureAlgorithm', signature_algorithm)

    # 模拟签名
    signature_value = univ.BitString.fromOctetString(os.urandom(256))
    certificate.setComponentByName('signature', signature_value)

    certificate_der = encoder.encode(certificate)
    return certificate_der


def create_pkcs7_signed_data(content, certificate_data, private_key_data, digest_algorithm='sha256'):
    """创建 PKCS#7 签名数据（未修改）"""
    # 代码同上，未修改
    # content_info = rfc2315.ContentInfo()
    # content_info['contentType'] = rfc2315.data
    #
    # content_data = rfc2315.Data()
    # content_data.setComponentByPosition(0, univ.OctetString(content))
    # content_info['content'] = content_data

    signed_data = rfc2315.SignedData()
    signed_data['version'] = univ.Integer(1)

    digest_algorithms = rfc2315.DigestAlgorithmIdentifiers()

    if digest_algorithm.lower() == 'sha1':
        digest_alg = rfc2315.DigestAlgorithmIdentifier()
        digest_alg['algorithm'] = rfc3279.id_sha1
        digest_algorithms.setComponentByPosition(0, digest_alg)
    elif digest_algorithm.lower() == 'sha256':
        digest_alg = rfc2315.DigestAlgorithmIdentifier()
        digest_alg['algorithm'] = rfc4055.id_sha256  # 使用 rfc4055 中的 id_sha256
        digest_algorithms.setComponentByPosition(0, digest_alg)
    else:
        raise ValueError(f"不支持的摘要算法: {digest_algorithm}")

    signed_data['digestAlgorithms'] = digest_algorithms

    certificates = rfc2315.ExtendedCertificatesAndCertificates()

    cert, _ = decoder.decode(certificate_data, asn1Spec=rfc2459.Certificate())
    certificates[0]['certificate']=cert

    signed_data['certificates'] = certificates.subtype(
            implicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatConstructed, 0))

    signer_infos = rfc2315.SignerInfos()
    signer_info = rfc2315.SignerInfo()
    signer_info['version'] = univ.Integer(1)

    issuer_and_serial = rfc2315.IssuerAndSerialNumber()
    issuer_and_serial['issuer'] = cert['tbsCertificate']['issuer']
    issuer_and_serial['serialNumber'] = cert['tbsCertificate']['serialNumber']
    signer_info['issuerAndSerialNumber'] = issuer_and_serial

    digest_alg = rfc2315.DigestAlgorithmIdentifier()
    if digest_algorithm.lower() == 'sha1':
        digest_alg['algorithm'] = rfc3279.id_sha1
    elif digest_algorithm.lower() == 'sha256':
        digest_alg['algorithm'] = rfc4055.id_sha256  # 使用 rfc4055 中的 id_sha256

    signer_info['digestAlgorithm'] = digest_alg

    if digest_algorithm.lower() == 'sha1':
        digest = hashlib.sha1(content).digest()
    elif digest_algorithm.lower() == 'sha256':
        digest = hashlib.sha256(content).digest()

    # 2. 定义 OID 常量
    OID_CONTENT_TYPE = univ.ObjectIdentifier('1.2.840.113549.1.9.3')
    OID_MESSAGE_DIGEST = univ.ObjectIdentifier('1.2.840.113549.1.9.4')
    OID_SIGNING_TIME = univ.ObjectIdentifier('1.2.840.113549.1.9.5')
    OID_DATA = univ.ObjectIdentifier('1.2.840.113549.1.7.1')  # PKCS#7 data

    # 3. 创建各个属性

    # 内容类型属性 (contentType)
    content_type_attr = rfc2315.Attribute()
    content_type_attr['type'] = OID_CONTENT_TYPE

    # 属性值 - SET OF AttributeValue
    attr_values = univ.SetOf()
    attr_values.append(OID_DATA)
    content_type_attr['values'] = attr_values

    # 消息摘要属性 (messageDigest)
    message_digest_attr = rfc2315.Attribute()
    message_digest_attr['type'] = OID_MESSAGE_DIGEST

    # 属性值 - SET OF AttributeValue
    attr_values = univ.SetOf()
    attr_values.append(univ.OctetString(digest))
    message_digest_attr['values'] = attr_values

    # 签名时间属性 (signingTime)
    signing_time_attr = rfc2315.Attribute()
    signing_time_attr['type'] = OID_SIGNING_TIME

    # 属性值 - SET OF AttributeValue
    attr_values = univ.SetOf()
    # 使用 UTCTime 格式 (RFC2315 要求)
    current_time = datetime.datetime.now().strftime('%y%m%d%H%M%SZ')
    attr_values.append(useful.UTCTime(current_time))
    signing_time_attr['values'] = attr_values

    # 4. 构建 AuthenticatedAttributes (SET OF Attribute)
    authenticated_attrs = univ.SetOf(
        componentType=rfc2315.Attribute()
    )

    # 添加属性 - 顺序很重要！
    authenticated_attrs.extend([
        content_type_attr,
        message_digest_attr,
        signing_time_attr
    ])

    signer_info['authenticatedAttributes'] = authenticated_attrs.subtype(
            implicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatConstructed, 0))

    digest_encryption_alg = rfc2315.DigestEncryptionAlgorithmIdentifier()
    digest_encryption_alg['algorithm'] = rfc3279.rsaEncryption
    signer_info['digestEncryptionAlgorithm'] = digest_encryption_alg

    encrypted_digest = univ.OctetString(os.urandom(256))
    signer_info['encryptedDigest'] = encrypted_digest

    signer_infos.setComponentByPosition(0, signer_info)
    signed_data['signerInfos'] = signer_infos

    final_content_info = rfc2315.ContentInfo()
    final_content_info['contentType'] = rfc2315.signedData

    final_content_info['content'] = signed_data

    der_encoded = encoder.encode(final_content_info)

    return der_encoded


def main():
    """主函数（未修改）"""
    try:
        print("开始生成 PKCS#7 签名数据...")

        print("正在生成 RSA 密钥对...")
        public_key_der, private_key_der = generate_rsa_key_pair(key_size=2048)

        print(f"公钥生成成功，大小: {len(public_key_der)} 字节")
        print(f"私钥生成成功，大小: {len(private_key_der)} 字节")

        print("正在生成 X.509 证书...")
        certificate_der = generate_certificate(public_key_der, validity_days=365)

        print(f"证书生成成功，大小: {len(certificate_der)} 字节")

        content = b"Hello, this is a sample content to be signed."
        print(f"要签名的内容: {content.decode('utf-8')}")

        print("正在创建 PKCS#7 签名数据...")
        pkcs7_data = create_pkcs7_signed_data(
            content=content,
            certificate_data=certificate_der,
            private_key_data=private_key_der,
            digest_algorithm='sha256'
        )

        with open('signed_data.der', 'wb') as f:
            f.write(pkcs7_data)

        print(f"PKCS#7 签名数据已创建并保存到 signed_data.der")
        print(f"数据大小: {len(pkcs7_data)} 字节")

        with open('certificate.der', 'wb') as f:
            f.write(certificate_der)

        with open('private_key.der', 'wb') as f:
            f.write(private_key_der)

        print(f"证书已保存到 certificate.der")
        print(f"私钥已保存到 private_key.der")

        base64_data = base64.b64encode(pkcs7_data).decode('utf-8')
        print("\nPKCS#7 签名数据的 Base64 编码:")
        print(base64_data[:100] + "..." + base64_data[-50:])

        print("\n操作完成!")

    except Exception as e:
        print(f"生成 PKCS#7 签名数据时出错: {str(e)}")
        print(traceback.print_exc())


if __name__ == "__main__":
    main()